package love.forte.suspendreversal.processor

import com.google.devtools.ksp.*
import com.google.devtools.ksp.processing.*
import com.google.devtools.ksp.symbol.*
import com.squareup.kotlinpoet.*
import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
import com.squareup.kotlinpoet.ksp.*
import love.forte.suspendreversal.annotations.SuspendReversal
import java.time.Instant

/**
 * [parentAnnotation] 和 [annotation] 应必然有一个不为 null
 *
 */
internal data class AnnotationAndClassDeclaration(
    val parentAnnotation: SuspendReversal?,
    val annotation: SuspendReversal?,
    val declaration: KSClassDeclaration
) {
    init {
        require(parentAnnotation != null || annotation != null)
    }

    val nearestAnnotation: SuspendReversal = (annotation ?: parentAnnotation)!!
}

@OptIn(KspExperimental::class)
private fun KSAnnotated.findReversalAnnotation(): SuspendReversal? {
    return getAnnotationsByType(SuspendReversal::class).firstOrNull()
}

internal val ReversalAnnotationName = SuspendReversal::class.java.name

internal val JvmThrowsClassName = ClassName("kotlin.jvm", "Throws")
internal val CommonThrowsClassName = ClassName("kotlin", "Throws")

internal const val FILE_EX_PREFIX = "suspendreversal.generated"

/**
 *
 * @author ForteScarlet
 */
internal class SuspendInterfaceReversalProcessor(private val environment: SymbolProcessorEnvironment) :
    SymbolProcessor {

    @OptIn(KspExperimental::class)
    override fun process(resolver: Resolver): List<KSAnnotated> {
        environment.logger.info("resolver: $resolver")
        environment.logger.info("platforms: ${environment.platforms}")

        val isJs = environment.platforms.any { it is JsPlatformInfo }
        val isJvm = environment.platforms.any { it is JvmPlatformInfo }

        val symbols =
            resolver.getSymbolsWithAnnotation(ReversalAnnotationName, inDepth = true)
                .flatMap {
                    val currAnnotation = it.findReversalAnnotation() ?: return@flatMap emptySequence()

                    it.getAnnotationsByType(SuspendReversal::class).firstOrNull()

                    when (it) {
                        is KSFile -> it.declarations.filterIsInstance<KSClassDeclaration>().map { declaration ->
                            val childAnnotation = declaration.findReversalAnnotation()
                            AnnotationAndClassDeclaration(currAnnotation, childAnnotation, declaration)
                        }

                        is KSClassDeclaration -> sequenceOf(AnnotationAndClassDeclaration(null, currAnnotation, it))
                        else -> emptySequence()
                    }
                }
                .filter {
                    // must be an interface or abstract class
                    when (it.declaration.classKind) {
                        ClassKind.INTERFACE -> true
                        ClassKind.CLASS -> {
                            if (!it.declaration.isAbstract()) {
                                environment.logger.warn(
                                    "@SuspendReversal can only be used on interfaces and abstract classes",
                                    it.declaration
                                )
                                false
                            } else {
                                true
                            }
                        }

                        ClassKind.ANNOTATION_CLASS -> false

                        else -> {
                            environment.logger.warn(
                                "@SuspendReversal can only be used on interfaces and abstract classes",
                                it.declaration
                            )
                            false
                        }
                    }
                }
                .toList()

        data class FileSpecBuilderInfo(val builder: FileSpec.Builder, val source: List<KSFile> = emptyList())
        data class Key(val p: String, val f: String)

        val fileBuilders = mutableMapOf<Key, FileSpecBuilderInfo>()

        symbols.forEach { symbol ->
            val containingFile = symbol.declaration.containingFile
            val generatedFileName = (containingFile?.fileName?.substringBeforeLast(".kt")
                ?: "") + "SuspendReversals.$FILE_EX_PREFIX"
            val packageName = symbol.declaration.packageName.asString()

            val fileBuilder = fileBuilders.computeIfAbsent(Key(packageName, generatedFileName)) {
                val builder = FileSpec.builder(
                    packageName = packageName,
                    fileName = generatedFileName,
                )

                builder.addFileComment("Auto generated by kotlin-suspend-interface-reversal at %L\n", Instant.now())
                builder.addFileComment("see https://github.com/ForteScarlet/kotlin-suspend-interface-reversal")

                FileSpecBuilderInfo(builder, symbol.declaration.containingFile?.let { listOf(it) } ?: emptyList())
            }
            // blocking type, async type

            val typeParameterResolver = symbol.declaration.typeParameters.toTypeParameterResolver()
            val superClassName = symbol.declaration.toClassName().let { name ->
                if (symbol.declaration.typeParameters.isNotEmpty()) {
                    name.parameterizedBy(symbol.declaration.typeParameters.map {
                        it.toTypeVariableName(
                            typeParameterResolver
                        )
                    })
                } else name
            }

            when {
                isJvm -> {
                    resolveJvm(
                        environment,
                        symbol,
                        fileBuilder.builder,
                        typeParameterResolver,
                        superClassName
                    )
                }

                isJs -> {
                    resolveJs(
                        environment,
                        symbol,
                        fileBuilder.builder,
                        typeParameterResolver,
                        superClassName
                    )
                }
            }

        }

        fileBuilders.values.forEach { (builder, sources) ->
            builder.build().writeTo(
                codeGenerator = environment.codeGenerator,
                dependencies = Dependencies(
                    aggregating = true,
                    sources = sources.toTypedArray()
                )
            )
        }

        return emptyList()
    }
}

internal fun generateReversalTypeSpecBuilder(
    typeParameterResolver: TypeParameterResolver,
    superClassName: TypeName,
    environment: SymbolProcessorEnvironment,
    symbol: AnnotationAndClassDeclaration,
    isEnabled: Boolean,
    classNamePrefix: String,
    classNameSuffix: String,
): TypeSpec.Builder? {
    val declaration = symbol.declaration
    val declarationName = declaration.simpleName.asString()

    val builder: TypeSpec.Builder = if (isEnabled) {
        when (declaration.classKind) {
            ClassKind.CLASS -> {
                val className = classNamePrefix + declaration.simpleName.asString() + classNameSuffix

                TypeSpec.classBuilder(className)
                    .apply {
                        superclass(superClassName)
                        // build primary constructor
                        val sourcePrimaryConstructor = declaration.primaryConstructor
                        if (sourcePrimaryConstructor != null) {
                            val parameters = sourcePrimaryConstructor.parameters.map { param ->
                                ParameterSpec.builder(
                                    param.name?.asString() ?: "",
                                    param.type.toTypeName(typeParameterResolver),
                                ).apply {
                                    if (param.isVararg) {
                                        addModifiers(KModifier.VARARG)
                                    }
                                }.build()
                            }

                            val primaryConstructor = FunSpec.constructorBuilder().apply {
                                addParameters(parameters)
                            }.build()

                            primaryConstructor(primaryConstructor)
                            addSuperclassConstructorParameter(
                                parameters.joinToString { "%N" },
                                args = (parameters as Collection<Any>).toTypedArray()
                            )
                        }
                    }
            }

            ClassKind.INTERFACE -> {
                environment.logger.info("INTERFACE declaration: $declaration")

                TypeSpec.interfaceBuilder(classNamePrefix + declaration.simpleName.asString())
                    .apply {
                        addSuperinterface(superClassName)
                    }
            }

            else -> return null
        }.apply {
            // modifiers
            addModifiers(declaration.modifiers.mapNotNull { it.toKModifier() })
            // type vs
            addTypeVariables(declaration.typeParameters.map { it.toTypeVariableName(typeParameterResolver) })
        }
    } else return null

    builder.addKdoc(
        """
        Generated suspend reversal type for [%N].
        @see %N
    """.trimIndent(), declarationName, declarationName
    )


    return builder
}

internal data class GeneratedReversalFunctions(
    val generatedFunction: FunSpec,
    val overriddenSuspendFunction: FunSpec
)

internal fun KSType.resolveClassDeclaration(): KSClassDeclaration? {
    return when (val declaration = declaration) {
        is KSClassDeclaration -> declaration
        is KSTypeAlias -> declaration.findActualType()
        else -> return null
    }
}

internal fun KSType.resolveClassDeclarationToClassName(): ClassName? {
    return resolveClassDeclaration()?.toClassName()
}

/**
 * 需要添加（拷贝）的注解们
 * - [JvmThrowsClassName]
 * - [CommonThrowsClassName]
 */
internal fun resolveIncludeAnnotations(builder: FunSpec.Builder, source: KSFunctionDeclaration) {
    source.annotations.filter {
        val anno = it.annotationType.resolve().resolveClassDeclarationToClassName() ?: return@filter false

        anno == JvmThrowsClassName
                || anno == CommonThrowsClassName
    }.forEach {
        builder.addAnnotation(it.toAnnotationSpec())
    }

}

/**
 * 粗略的判断是否需要标记继承
 */
internal fun shouldOverride(
    declaration: KSClassDeclaration,
    name: String,
    extensionReceiver: KSTypeReference?,
    params: List<KSValueParameter>
): Boolean {
    return declaration.getAllFunctions().any { func ->
        if (func.simpleName.asString() != name) {
            return@any false
        }

        val funcReceiver = func.extensionReceiver
        if (funcReceiver != extensionReceiver) return@any false
        val funParams = func.parameters
        if (funParams.size != params.size) return@any false

        for ((index, funParam) in funParams.withIndex()) {
            val param = params[index]
            if (param.type == funParam.type) continue

            val funParamClassDecl = funParam.type.resolve().resolveClassDeclaration()
            val paramClassDecl = param.type.resolve().resolveClassDeclaration()

            // skip.
            if (funParamClassDecl == null || paramClassDecl == null) return false
        }

        return true
    }
}
